import os
import csv
from tqdm import tqdm
import time
import math
import random
import numpy as np
random.seed(2024)
# from lightgbm import LGBMClassifier, Booster

from sklearn.metrics import roc_auc_score
# from sklearn.neural_network import MLPClassifier
from sklearn.tree import DecisionTreeClassifier
# from sklearn.svm import SVC

project_dir = os.path.dirname(os.path.abspath(__file__))
project_dir = project_dir.replace('scripts','')
import sys
sys.path.append(project_dir)

filename = os.path.join(project_dir, 'data/train_prob.csv')
train_features,train_labels = [],[]
test_features,test_labels = [],[]
with open(filename, 'r', encoding='utf-8') as f:
    lines = csv.reader(f)
    next(lines)
    for line in lines:
        if random.random() < 0.9:
            train_features.append([-math.log(float(i),10) for i in line[1:]])
            train_labels.append(int(line[0]))
        else:
            test_features.append([-math.log(float(i),10) for i in line[1:]])
            test_labels.append(int(line[0]))

train_features = np.array(train_features)
train_labels = np.array(train_labels)
test_features = np.array(test_features)
test_labels = np.array(test_labels)

# 构建决策树分类器
clf = DecisionTreeClassifier()
# clf = LGBMClassifier(
#         max_depth=3,
#         learning_rate=0.1,
#         n_estimators=200,
#         objective='multiclass',
#         num_class=2,
#         booster='gbtree',
#         min_child_weight=2,
#         subsample=0.8,
#         colsample_bytree=0.8,
#         reg_alpha=0,
#         reg_lambda=1,
#         seed=0)
clf.fit(train_features, train_labels)

# 在测试集上进行预测
y_pred = clf.predict(test_features)

# 计算AUC分数
auc = roc_auc_score(test_labels, y_pred)
print("AUC分数: ", auc)